# -*-coding:utf-8-*-
"""
Created on 2021.7.15

@author:夜剑听雨
"""

import numpy as np
import matplotlib.pyplot as plt
import segyio
# pip install segyio
# 读取SEG-Y格式地震数据
def ReadSEGYData(filename):
    print("### Reading SEGY-formatted Seismic Data:")
    print("Data file-->[%s]" %(filename))

    with segyio.open(filename, "r", ignore_geometry=True)as f:
        f.mmap()
        nTrace = f.tracecount
        nSample = f.bin[segyio.BinField.Samples]
        startT = 0
        deltaT = f.bin[segyio.BinField.Interval]
        print("     Number of Trace   = %d" % (nTrace))
        print("     Number of Samples = %d" % (nSample))
        print("     Start Samples     = %d" % (startT))
        print("     Sampling Rate     = %d" % (deltaT))
        data2D = np.asarray([np.copy(x) for x in f.trace[:]]).T
    my_data = [data2D, nTrace, nSample, startT, deltaT]
    f.close()
    return my_data
# 加载文件夹数据
data2D = "./test.sgy"
my_data = ReadSEGYData(data2D)
data2D = my_data[0]
print(data2D.shape)
# 直接对原始数据可视化
clip = 3   # 显示范围，负值越大越明显
vmin, vmax = -clip, clip

# 彩色图像# 纵横比为 0.5   29760
plt.figure('1')
plt.imshow(data2D[0:my_data[1], 0: 1000], cmap=plt.cm.seismic,
           interpolation='nearest', aspect=1, vmin=vmin, vmax=vmax)
plt.xlabel('Trace number')
plt.ylabel('Time sample')
# 灰色图像
plt.figure('2')
plt.imshow(data2D[0:my_data[1], 0:1000], cmap=plt.cm.bone,
           interpolation='nearest', aspect=0.5, vmin=vmin, vmax=vmax)
plt.xlabel('Trace number')
plt.ylabel('Time sample')
# 灰度图像
plt.figure('3')
plt.imshow(data2D[0:my_data[1], 0:1000], cmap='Greys',
           interpolation='nearest', aspect=2, vmin=vmin, vmax=vmax)
plt.xlabel('Trace number')
plt.ylabel('Time sample')
plt.show()
